package main

import (
	"bytes"
	"fmt"
	"io"
	log "log/slog"
	"net"
	"net/url"
	"strings"
)

const (
	MethodConnect = "CONNECT"
	BufSize       = 64
)

type connection struct {
	conn net.Conn
}

func (c *connection) serve() {
	defer c.conn.Close()

	log.Info(fmt.Sprintf("Receiving request from %v", c.conn.RemoteAddr()))
	data, remoteAddr, secure, err := readReqLine(c.conn)
	if err != nil {
		log.Error(fmt.Sprintf("Unable to read request info from %v: %v", c.conn.RemoteAddr(), err))
	}

	log.Info(fmt.Sprintf("Establishing connection to remote %v", remoteAddr))
	remoteConn, err := net.Dial("tcp", remoteAddr)
	if err != nil {
		log.Error(fmt.Sprintf("Failed to connect to host %s: %v", remoteAddr, err))
		return
	}

	defer remoteConn.Close()

	if secure {
		c.conn.Write([]byte("HTTP/1.1 200 Connection Established\r\n\r\n"))
	} else {
		_, err := remoteConn.Write(data)
		if err != nil {
			log.Error(fmt.Sprintf("Failed to write request header to remote host! %v", err.Error()))
			return
		}
	}

	log.Info(fmt.Sprintf("Start tunnel connection %v <-> %v", c.conn.RemoteAddr(), remoteAddr))
	defer log.Info(fmt.Sprintf("Finish tunnel connection %v <-> %v", c.conn.RemoteAddr(), remoteAddr))

	go io.Copy(remoteConn, c.conn)
	io.Copy(c.conn, remoteConn)
}

func readReqLine(rc net.Conn) (data []byte, addr string, secureRequest bool, err error) {
	reqFirstLine, data, err := retrieveReqFirstLine(rc)
	if err != nil {
		log.Error(fmt.Sprintf("Failed to read request info from request connnection: %v", err.Error()))
		return
	}

	method, uri, _, ok := parseReqLine(reqFirstLine)

	if !ok {
		log.Error("Malformed request")
		return
	}

	if method == MethodConnect {
		secureRequest = true
		data, err = drainConnectReqHeader(rc, data)
		cherr(err)
	}

	if secureRequest {
		uri = "https://" + uri
	}
	u, err := url.Parse(uri)
	if err != nil {
		log.Error(fmt.Sprintf("Failed to parse request uri %s. error: %v", uri, err))
		return
	}

	if secureRequest {
		addr = u.Host
	} else {
		addr = u.Host
		if !strings.Contains(addr, ":") {
			addr += ":80"
		}
	}
	return data, addr, secureRequest, nil
}

func retrieveReqFirstLine(rc net.Conn) (reqFirstLine string, data []byte, err error) {
	buf := make([]byte, BufSize)

	for {
		var bytesRead int
		bytesRead, err = rc.Read(buf)
		cherr(err)

		data = append(data, buf[:bytesRead]...)

		index := bytes.Index(data, []byte("\r\n"))

		if index != -1 {
			reqFirstLine = string(data[:index])
			break
		}
	}
	return reqFirstLine, data, nil
}

func parseReqLine(line string) (method, path, ver string, ok bool) {
	tokens := strings.Split(line, " ")
	if len(tokens) != 3 {
		return
	}
	return tokens[0], tokens[1], tokens[2], true
}

func drainConnectReqHeader(rc net.Conn, data []byte) ([]byte, error) {
	for {
		buf := make([]byte, BufSize)

		n, err := rc.Read(buf)
		cherr(err)

		data = append(data, buf[:n]...)
		if n < BufSize || buf[BufSize-1] == byte('\n') {
			// log.Info("Drained https connnect request header")
			break
		}
	}
	return data, nil
}

func newConnection(conn net.Conn) *connection {
	return &connection{
		conn: conn,
	}
}

func cherr(err error) {
	if err != nil {
		panic(err)
	}
}
